# models.py
from pydantic import BaseModel, Field
from typing import List, Dict, Any, Optional, Union

# --- 输入模型 ---

class Position(BaseModel):
    lon: float
    lat: float
    alt: float

class UAVStart(BaseModel):
    uav_id: int
    position: Position

class WaypointGoal(BaseModel):
    uav_id: int
    position: Position

class Task(BaseModel):
    priority: int
    taskType: str = Field(..., description="任务类型，例如 'patrol' 或 'waypoint'")
    time_constraint: float
    uav_ids: List[int]
    start: List[UAVStart]
    vertices: Optional[List[Position]] = None
    goals: Optional[List[WaypointGoal]] = None

class PlannerConfig(BaseModel):
    """
    规划器相关参数，由终端部提供
    """
    ref_lat: float = Field(29.35541411, description="参考点纬度")
    ref_lon: float = Field(106.0073147, description="参考点经度")
    robot_radius: float = Field(2.0, description="半径")
    coarse_grid_size: float = Field(20.0, description="粗规划网格大小")
    fine_grid_size: float = Field(5.0, description="精细规划网格大小")
    # 示例：添加Astar地图边界
    map_bounds: List[float] = Field([-300, 0, 400, 400], description="地图边界 [min_x, min_y, max_x, max_y]")

class PlanningRequest(BaseModel):
    """
    API的完整请求体
    """
    tasks: List[Task]
    config: PlannerConfig

# --- 输出模型 ---

class TrajectoryData(BaseModel):
    trajectory: List[Position]
    suggested_planning_speed: Optional[float] = None
    suggested_patrol_speed: Optional[float] = None

class PlanningResponseData(BaseModel):
    trajectories: Dict[str, TrajectoryData]
    time_paths: Dict[str, List[List[Union[int, float]]]]

class PlanningResponse(BaseModel):
    status: str
    message: str
    data: Optional[PlanningResponseData] = None
